"""

inference on SR-GAN super resolution model
this will load a blurry image as input
and output a more detailed image as output

"""
from alfred.dl.tf.common import mute_tf
mute_tf()
import tensorflow as tf
from model.srgan import generator, discriminator
from model import resolve_single
from utils import load_image
import os
import matplotlib.pyplot as plt
import cv2
import numpy as np


weights_dir = 'weights/srgan'
weights_file = lambda filename: os.path.join(weights_dir, filename)
save_dir = 'results/'
os.makedirs(save_dir, exist_ok=True)
target_shape = (118, 124) # h, w


def resolve_and_plot(lr_image_path):
    lr = load_image(lr_image_path)
    
    pre_sr = resolve_single(pre_generator, lr)
    gan_sr = resolve_single(gan_generator, lr)

    pre_sr = pre_sr.numpy()
    gan_sr = gan_sr.numpy()
    
    cv2.imshow('low resolution', lr)
    cv2.imshow('pre super', pre_sr)
    cv2.imshow('gan super', gan_sr)

    cv2.imwrite(os.path.join(save_dir, '{}_ori.png'.format(os.path.basename(lr_image_path))), lr)
    cv2.imwrite(os.path.join(save_dir, '{}_pre_sr.png'.format(os.path.basename(lr_image_path))), pre_sr)
    cv2.imwrite(os.path.join(save_dir, '{}_gan_sr.png'.format(os.path.basename(lr_image_path))), gan_sr)
    cv2.waitKey(0)


def resolve_and_plot2(lr_image_path):
    lr = load_image(lr_image_path)
    if lr.shape[0] != target_shape[0] and lr.shape[1] != target_shape[1]:
        # reshape image
        print('resize input: {} vs {}'.format(lr.shape, target_shape))
        lr = cv2.resize(lr, (target_shape[1], target_shape[0]))
        print('after: {}'.format(lr.shape))
        if lr.shape[-1] != 3:
            lr = lr[:, :, :-1]
            print('trim a channel: {}'.format(lr.shape))
    
    pre_sr = resolve_single(pre_generator, lr)
    gan_sr = resolve_single(gan_generator, lr)

    pre_sr = pre_sr.numpy()
    gan_sr = gan_sr.numpy()
    
    print('resize lr: ', gan_sr.shape)
    lr_larger = cv2.resize(lr, (gan_sr.shape[1], gan_sr.shape[0]))
    cv2.putText(gan_sr, 'Super Resolution', (20, 20), cv2.FONT_HERSHEY_COMPLEX, 0.6, (255, 0, 255), 2)
    cv2.putText(lr_larger, 'Original', (20, 20), cv2.FONT_HERSHEY_COMPLEX, 0.6, (255, 0, 255), 2)

    compare_img = np.concatenate((lr_larger, gan_sr), axis=1)
    cv2.imshow('low resolution', lr_larger)
    cv2.imshow('gan super', gan_sr)
    cv2.imshow('compare', compare_img)

    cv2.imwrite(os.path.join(save_dir, '{}_compare.png'.format(os.path.basename(lr_image_path))), compare_img)
    cv2.waitKey(0)


if __name__ == "__main__":
    pre_generator = generator()
    gan_generator = generator()

    pre_generator.load_weights(weights_file('pre_generator.h5'))
    gan_generator.load_weights(weights_file('gan_generator.h5'))

    print('model loaded.')

    resolve_and_plot2('images/0869x4-crop.png')
    resolve_and_plot2('images/0829x4-crop.png')
    resolve_and_plot2('images/0851x4-crop.png')
    resolve_and_plot2('images/1.png')